{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Agentic RAG with Autogen\n",
    "\n",
    "This notebook demonstrates implementing Retrieval-Augmented Generation (RAG) using Autogen agents with enhanced evaluation capabilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: chromadb in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (0.6.3)\n",
      "Requirement already satisfied: build>=1.0.3 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (1.2.2.post1)\n",
      "Requirement already satisfied: pydantic>=1.9 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (2.10.6)\n",
      "Requirement already satisfied: chroma-hnswlib==0.7.6 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (0.7.6)\n",
      "Requirement already satisfied: fastapi>=0.95.2 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (0.115.8)\n",
      "Requirement already satisfied: uvicorn>=0.18.3 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from uvicorn[standard]>=0.18.3->chromadb) (0.34.0)\n",
      "Requirement already satisfied: numpy>=1.22.5 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (2.2.3)\n",
      "Requirement already satisfied: posthog>=2.4.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (3.14.1)\n",
      "Requirement already satisfied: typing_extensions>=4.5.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (4.12.2)\n",
      "Requirement already satisfied: onnxruntime>=1.14.1 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (1.20.1)\n",
      "Requirement already satisfied: opentelemetry-api>=1.2.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (1.30.0)\n",
      "Requirement already satisfied: opentelemetry-exporter-otlp-proto-grpc>=1.2.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (1.30.0)\n",
      "Requirement already satisfied: opentelemetry-instrumentation-fastapi>=0.41b0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (0.51b0)\n",
      "Requirement already satisfied: opentelemetry-sdk>=1.2.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (1.30.0)\n",
      "Requirement already satisfied: tokenizers>=0.13.2 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (0.21.0)\n",
      "Requirement already satisfied: pypika>=0.48.9 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (0.48.9)\n",
      "Requirement already satisfied: tqdm>=4.65.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (4.67.1)\n",
      "Requirement already satisfied: overrides>=7.3.1 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (7.7.0)\n",
      "Requirement already satisfied: importlib-resources in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (6.5.2)\n",
      "Requirement already satisfied: grpcio>=1.58.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (1.70.0)\n",
      "Requirement already satisfied: bcrypt>=4.0.1 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (4.2.1)\n",
      "Requirement already satisfied: typer>=0.9.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (0.15.1)\n",
      "Requirement already satisfied: kubernetes>=28.1.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (32.0.0)\n",
      "Requirement already satisfied: tenacity>=8.2.3 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (9.0.0)\n",
      "Requirement already satisfied: PyYAML>=6.0.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (6.0.2)\n",
      "Requirement already satisfied: mmh3>=4.0.1 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (5.1.0)\n",
      "Requirement already satisfied: orjson>=3.9.12 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (3.10.15)\n",
      "Requirement already satisfied: httpx>=0.27.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (0.28.1)\n",
      "Requirement already satisfied: rich>=10.11.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from chromadb) (13.9.4)\n",
      "Requirement already satisfied: packaging>=19.1 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from build>=1.0.3->chromadb) (24.2)\n",
      "Requirement already satisfied: pyproject_hooks in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from build>=1.0.3->chromadb) (1.2.0)\n",
      "Requirement already satisfied: colorama in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from build>=1.0.3->chromadb) (0.4.6)\n",
      "Requirement already satisfied: starlette<0.46.0,>=0.40.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from fastapi>=0.95.2->chromadb) (0.45.3)\n",
      "Requirement already satisfied: anyio in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from httpx>=0.27.0->chromadb) (4.8.0)\n",
      "Requirement already satisfied: certifi in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from httpx>=0.27.0->chromadb) (2025.1.31)\n",
      "Requirement already satisfied: httpcore==1.* in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from httpx>=0.27.0->chromadb) (1.0.7)\n",
      "Requirement already satisfied: idna in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from httpx>=0.27.0->chromadb) (3.10)\n",
      "Requirement already satisfied: h11<0.15,>=0.13 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from httpcore==1.*->httpx>=0.27.0->chromadb) (0.14.0)\n",
      "Requirement already satisfied: six>=1.9.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from kubernetes>=28.1.0->chromadb) (1.17.0)\n",
      "Requirement already satisfied: python-dateutil>=2.5.3 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from kubernetes>=28.1.0->chromadb) (2.9.0.post0)\n",
      "Requirement already satisfied: google-auth>=1.0.1 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from kubernetes>=28.1.0->chromadb) (2.38.0)\n",
      "Requirement already satisfied: websocket-client!=0.40.0,!=0.41.*,!=0.42.*,>=0.32.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from kubernetes>=28.1.0->chromadb) (1.8.0)\n",
      "Requirement already satisfied: requests in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from kubernetes>=28.1.0->chromadb) (2.32.3)\n",
      "Requirement already satisfied: requests-oauthlib in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from kubernetes>=28.1.0->chromadb) (2.0.0)\n",
      "Requirement already satisfied: oauthlib>=3.2.2 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from kubernetes>=28.1.0->chromadb) (3.2.2)\n",
      "Requirement already satisfied: urllib3>=1.24.2 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from kubernetes>=28.1.0->chromadb) (2.3.0)\n",
      "Requirement already satisfied: durationpy>=0.7 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from kubernetes>=28.1.0->chromadb) (0.9)\n",
      "Requirement already satisfied: coloredlogs in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from onnxruntime>=1.14.1->chromadb) (15.0.1)\n",
      "Requirement already satisfied: flatbuffers in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from onnxruntime>=1.14.1->chromadb) (25.2.10)\n",
      "Requirement already satisfied: protobuf in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from onnxruntime>=1.14.1->chromadb) (5.29.3)\n",
      "Requirement already satisfied: sympy in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from onnxruntime>=1.14.1->chromadb) (1.13.3)\n",
      "Requirement already satisfied: deprecated>=1.2.6 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from opentelemetry-api>=1.2.0->chromadb) (1.2.18)\n",
      "Requirement already satisfied: importlib-metadata<=8.5.0,>=6.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from opentelemetry-api>=1.2.0->chromadb) (8.5.0)\n",
      "Requirement already satisfied: googleapis-common-protos~=1.52 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from opentelemetry-exporter-otlp-proto-grpc>=1.2.0->chromadb) (1.67.0)\n",
      "Requirement already satisfied: opentelemetry-exporter-otlp-proto-common==1.30.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from opentelemetry-exporter-otlp-proto-grpc>=1.2.0->chromadb) (1.30.0)\n",
      "Requirement already satisfied: opentelemetry-proto==1.30.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from opentelemetry-exporter-otlp-proto-grpc>=1.2.0->chromadb) (1.30.0)\n",
      "Requirement already satisfied: opentelemetry-instrumentation-asgi==0.51b0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (0.51b0)\n",
      "Requirement already satisfied: opentelemetry-instrumentation==0.51b0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (0.51b0)\n",
      "Requirement already satisfied: opentelemetry-semantic-conventions==0.51b0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (0.51b0)\n",
      "Requirement already satisfied: opentelemetry-util-http==0.51b0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (0.51b0)\n",
      "Requirement already satisfied: wrapt<2.0.0,>=1.0.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from opentelemetry-instrumentation==0.51b0->opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (1.17.2)\n",
      "Requirement already satisfied: asgiref~=3.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from opentelemetry-instrumentation-asgi==0.51b0->opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (3.8.1)\n",
      "Requirement already satisfied: monotonic>=1.5 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from posthog>=2.4.0->chromadb) (1.6)\n",
      "Requirement already satisfied: backoff>=1.10.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from posthog>=2.4.0->chromadb) (2.2.1)\n",
      "Requirement already satisfied: annotated-types>=0.6.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from pydantic>=1.9->chromadb) (0.7.0)\n",
      "Requirement already satisfied: pydantic-core==2.27.2 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from pydantic>=1.9->chromadb) (2.27.2)\n",
      "Requirement already satisfied: markdown-it-py>=2.2.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from rich>=10.11.0->chromadb) (3.0.0)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from rich>=10.11.0->chromadb) (2.19.1)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.16.4 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from tokenizers>=0.13.2->chromadb) (0.28.1)\n",
      "Requirement already satisfied: click>=8.0.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from typer>=0.9.0->chromadb) (8.1.8)\n",
      "Requirement already satisfied: shellingham>=1.3.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from typer>=0.9.0->chromadb) (1.5.4)\n",
      "Requirement already satisfied: httptools>=0.6.3 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from uvicorn[standard]>=0.18.3->chromadb) (0.6.4)\n",
      "Requirement already satisfied: python-dotenv>=0.13 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from uvicorn[standard]>=0.18.3->chromadb) (1.0.1)\n",
      "Requirement already satisfied: watchfiles>=0.13 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from uvicorn[standard]>=0.18.3->chromadb) (1.0.4)\n",
      "Requirement already satisfied: websockets>=10.4 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from uvicorn[standard]>=0.18.3->chromadb) (15.0)\n",
      "Requirement already satisfied: cachetools<6.0,>=2.0.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from google-auth>=1.0.1->kubernetes>=28.1.0->chromadb) (5.5.1)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from google-auth>=1.0.1->kubernetes>=28.1.0->chromadb) (0.4.1)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from google-auth>=1.0.1->kubernetes>=28.1.0->chromadb) (4.9)\n",
      "Requirement already satisfied: filelock in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from huggingface-hub<1.0,>=0.16.4->tokenizers>=0.13.2->chromadb) (3.17.0)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from huggingface-hub<1.0,>=0.16.4->tokenizers>=0.13.2->chromadb) (2025.2.0)\n",
      "Requirement already satisfied: zipp>=3.20 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from importlib-metadata<=8.5.0,>=6.0->opentelemetry-api>=1.2.0->chromadb) (3.21.0)\n",
      "Requirement already satisfied: mdurl~=0.1 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from markdown-it-py>=2.2.0->rich>=10.11.0->chromadb) (0.1.2)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from requests->kubernetes>=28.1.0->chromadb) (3.4.1)\n",
      "Requirement already satisfied: sniffio>=1.1 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from anyio->httpx>=0.27.0->chromadb) (1.3.1)\n",
      "Requirement already satisfied: humanfriendly>=9.1 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from coloredlogs->onnxruntime>=1.14.1->chromadb) (10.0)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from sympy->onnxruntime>=1.14.1->chromadb) (1.3.0)\n",
      "Requirement already satisfied: pyreadline3 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from humanfriendly>=9.1->coloredlogs->onnxruntime>=1.14.1->chromadb) (3.5.4)\n",
      "Requirement already satisfied: pyasn1<0.7.0,>=0.4.6 in e:\\coding material\\python pr\\mlsa blog\\github\\ai-agents-for-beginners\\venv\\lib\\site-packages (from pyasn1-modules>=0.2.1->google-auth>=1.0.1->kubernetes>=28.1.0->chromadb) (0.6.1)\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "[notice] A new release of pip is available: 24.0 -> 25.0.1\n",
      "[notice] To update, run: python.exe -m pip install --upgrade pip\n"
     ]
    }
   ],
   "source": [
    "# Make sure to run this cell before running the rest of the notebook\n",
    "%pip install chromadb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SQLite Version Fix\n",
    "If you encounter the error:\n",
    "```\n",
    "RuntimeError: Your system has an unsupported version of sqlite3. Chroma requires sqlite3 >= 3.35.0\n",
    "```\n",
    "\n",
    "Uncomment this code block at the start of your notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %pip install pysqlite3-binary\n",
    "# __import__('pysqlite3')\n",
    "# import sys\n",
    "# sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import asyncio\n",
    "from typing import List, Dict\n",
    "\n",
    "from autogen_agentchat.agents import AssistantAgent\n",
    "from autogen_core import CancellationToken\n",
    "from autogen_agentchat.messages import TextMessage\n",
    "from azure.core.credentials import AzureKeyCredential\n",
    "from autogen_ext.models.azure import AzureAIChatCompletionClient\n",
    "\n",
    "import chromadb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the Client \n",
    "\n",
    "First, we initialize the Azure AI Chat Completion Client. This client will be used to interact with the Azure OpenAI service to generate responses to user queries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = AzureAIChatCompletionClient(\n",
    "    model=\"gpt-4o-mini\",\n",
    "    endpoint=\"https://models.inference.ai.azure.com\",\n",
    "    credential=AzureKeyCredential(os.environ[\"GITHUB_TOKEN\"]),\n",
    "    model_info={\n",
    "        \"json_output\": True,\n",
    "        \"function_calling\": True,\n",
    "        \"vision\": True,\n",
    "        \"family\": \"unknown\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vector Database Initialization\n",
    "\n",
    "We initialize ChromaDB with persistent storage and add enhanced sample documents. ChromaDB will be used to store and retrieve documents that provide context for generating accurate responses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize ChromaDB with persistent storage\n",
    "chroma_client = chromadb.PersistentClient(path=\"./chroma_db\")\n",
    "collection = chroma_client.create_collection(\n",
    "    name=\"travel_documents\",\n",
    "    metadata={\"description\": \"travel_service\"},\n",
    "    get_or_create=True\n",
    ")\n",
    "\n",
    "# Enhanced sample documents\n",
    "documents = [\n",
    "    \"Contoso Travel offers luxury vacation packages to exotic destinations worldwide.\",\n",
    "    \"Our premium travel services include personalized itinerary planning and 24/7 concierge support.\",\n",
    "    \"Contoso's travel insurance covers medical emergencies, trip cancellations, and lost baggage.\",\n",
    "    \"Popular destinations include the Maldives, Swiss Alps, and African safaris.\",\n",
    "    \"Contoso Travel provides exclusive access to boutique hotels and private guided tours.\"\n",
    "]\n",
    "\n",
    "# Add documents with metadata\n",
    "collection.add(\n",
    "    documents=documents,\n",
    "    ids=[f\"doc_{i}\" for i in range(len(documents))],\n",
    "    metadatas=[{\"source\": \"training\", \"type\": \"explanation\"} for _ in documents]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_retrieval_context(query: str) -> str:\n",
    "    results = collection.query(\n",
    "        query_texts=[query],\n",
    "        include=[\"documents\", \"metadatas\"],\n",
    "        n_results=2  # retrieve top 2 results for extra context\n",
    "    )\n",
    "    context_strings = []\n",
    "    if results and results.get(\"documents\") and len(results[\"documents\"][0]) > 0:\n",
    "        for doc, meta in zip(results[\"documents\"][0], results[\"metadatas\"][0]):\n",
    "            context_strings.append(f\"Document: {doc}\\nMetadata: {meta}\")\n",
    "    return \"\\n\\n\".join(context_strings) if context_strings else \"No results found\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Agent Configuration\n",
    "\n",
    "We configure the retrieval and assistant agents. The retrieval agent is specialized in finding relevant information using semantic search, while the assistant generates detailed responses based on the retrieved information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create agents with enhanced capabilities\n",
    "assistant = AssistantAgent(\n",
    "    name=\"assistant\",\n",
    "    model_client=client,\n",
    "    system_message=(\n",
    "        \"You are a helpful AI assistant that provides answers using ONLY the provided context. \"\n",
    "        \"Do NOT include any external information. Base your answer entirely on the context given below.\"\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RAGEvaluator Class\n",
    "\n",
    "We define the `RAGEvaluator` class to evaluate the response based on various metrics like response length, source citations, response time, and context relevance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RAGEvaluator:\n",
    "    def __init__(self):\n",
    "        self.responses: List[Dict] = []\n",
    "\n",
    "    def evaluate_response(self, query: str, response: str, context: List[str]) -> Dict:\n",
    "        # Basic metrics: response length, citation count, and a simple relevance score.\n",
    "        start_time = time.time()\n",
    "        metrics = {\n",
    "            'response_length': len(response),\n",
    "            'source_citations': sum(1 for doc in context if doc in response),\n",
    "            'evaluation_time': time.time() - start_time,\n",
    "            'context_relevance': self._calculate_relevance(query, context)\n",
    "        }\n",
    "        self.responses.append({\n",
    "            'query': query,\n",
    "            'response': response,\n",
    "            'metrics': metrics\n",
    "        })\n",
    "        return metrics\n",
    "\n",
    "    def _calculate_relevance(self, query: str, context: List[str]) -> float:\n",
    "        # Simple relevance score: fraction of the documents where the query appears.\n",
    "        return sum(1 for c in context if query.lower() in c.lower()) / len(context)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Query Processing with RAG\n",
    "\n",
    "We define the `ask_rag` function to send the query to the assistant, process the response, and evaluate it. This function handles the interaction with the assistant and uses the evaluator to measure the quality of the response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def ask_rag(query: str, evaluator: RAGEvaluator):\n",
    "    try:\n",
    "        retrieval_context = get_retrieval_context(query)\n",
    "        # Augment the query with the retrieval context.\n",
    "        augmented_query = (\n",
    "            f\"Retrieved Context:\\n{retrieval_context}\\n\\n\"\n",
    "            f\"User Query: {query}\\n\\n\"\n",
    "            \"Based ONLY on the above context, please provide the answer.\"\n",
    "        )\n",
    "\n",
    "        # Send the augmented query as a user message.\n",
    "        start_time = time.time()\n",
    "        response = await assistant.on_messages(\n",
    "            [TextMessage(content=augmented_query, source=\"user\")],\n",
    "            cancellation_token=CancellationToken(),\n",
    "        )\n",
    "        processing_time = time.time() - start_time\n",
    "\n",
    "        # Evaluate the response against our vector-store documents.\n",
    "        metrics = evaluator.evaluate_response(\n",
    "            query=query,\n",
    "            response=response.chat_message.content,\n",
    "            context=documents\n",
    "        )\n",
    "        return {\n",
    "            'response': response.chat_message.content,\n",
    "            'processing_time': processing_time,\n",
    "            'metrics': metrics,\n",
    "        }\n",
    "    except Exception as e:\n",
    "        print(f\"Error processing query: {e}\")\n",
    "        return None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example usage\n",
    "\n",
    "We initialize the evaluator and define the queries that we want to process and evaluate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def main():\n",
    "    evaluator = RAGEvaluator()\n",
    "    queries = [\n",
    "        \"Can you explain Contoso's travel insurance coverage?\", # Relevant document\n",
    "        \"What is Neural Network?\" # No relevant document\n",
    "    ]\n",
    "    for query in queries:\n",
    "        print(f\"\\nProcessing Query: {query}\")\n",
    "        result = await ask_rag(query, evaluator)\n",
    "        if result:\n",
    "            print(\"Response:\", result['response'])\n",
    "        print(\"\\n\" + \"=\"*60 + \"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the Script\n",
    "\n",
    "We check if the script is running in an interactive environment or a standard script, and run the main function accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Processing Query: Can you explain Contoso's travel insurance coverage?\n",
      "Response: Contoso's travel insurance covers medical emergencies, trip cancellations, and lost baggage.\n",
      "\n",
      "============================================================\n",
      "\n",
      "\n",
      "Processing Query: What is Neural Network?\n",
      "Response: The provided context does not contain any information regarding Neural Networks.\n",
      "\n",
      "============================================================\n",
      "\n"
     ]
    }
   ],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    if asyncio.get_event_loop().is_running():\n",
    "        await main()\n",
    "    else:\n",
    "        asyncio.run(main())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
